import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
from args import get_args
import os
from matplotlib.ticker import MultipleLocator


# 根据输入的两个loss列表画出曲线图
def plot_loss_curves(loss_list: np.ndarray,
                     titles: list = None,
                     xlabel: list = None,
                     ylabel: list = None,
                     labels: list = None,
                     colors: list = None,
                     batch_num=1,
                     spacing=3,
                     row_num=3):
    if len(loss_list.shape) < 3:
        raise ValueError('the shape of loss_list should be of 3')
    shape = loss_list.shape
    if not colors:
        colors = np.random.rand(shape[1], 3)
    if not labels:
        labels = ['train', 'test']
    if not titles:
        titles = [None] * shape[0]
    if not xlabel:
        xlabel = [None] * shape[0]
    if not ylabel:
        ylabel = [None] * shape[0]

    major_locator = MultipleLocator(spacing)
    epochs = shape[-1] // batch_num
    x = np.linspace(1, epochs, shape[-1])
    # 创建一个新的Figure对象
    fig = plt.figure()
    # 创建一个Axes对象
    for i in range(shape[0]):
        ax = fig.add_subplot(row_num, int(np.ceil(shape[0] / row_num)), i + 1)
        for index, loss in enumerate(loss_list[i]):
            ax.plot(x, loss, linewidth=2, color=colors[index], label=labels[index])
        ax.set_title(titles[i])
        ax.set_xlabel(xlabel[i])
        ax.set_ylabel(ylabel[i])
        ax.xaxis.set_major_locator(major_locator)
        ax.legend()
    plt.tight_layout()
    # plt.xlim(0.3, epochs + 0.5)
    # plt.show()
    plt.savefig(r'test.jpg')


def write_data(path, loss):
    df = pd.DataFrame(loss, columns=['loss'])
    writer = pd.ExcelWriter(path)
    df.to_excel(writer, index=False)
    writer.save()
    print('save successfully!')


def read_data(paths: list):
    size = len(paths)
    data = np.asarray(pd.read_excel(paths[0])).flatten()
    length = len(data)
    res = np.zeros(shape=(size, length))
    res[0] = np.asarray(data)
    for index, path in enumerate(paths[1:]):
        res[index + 1] = np.asarray(pd.read_excel(path)).flatten()
    return res


def get_loss(args):
    loss = np.zeros(shape=[len(args.models) * len(args.datasets), 2, 100])
    k = 0
    for _, model in enumerate(args.models):
        for _, dataset in enumerate(args.datasets):
            loss[k] = read_data(
                [os.path.join(args.loss_path, model, dataset, 'fake_train_loss.xlsx'),
                 os.path.join(args.loss_path, model, dataset, 'fake_test_loss.xlsx')])
            k += 1
    return loss


if __name__ == '__main__':
    args = get_args()

    font_size = 16
    dpi = 100
    params = {'axes.titlesize': font_size,
              'legend.fontsize': font_size,
              'figure.dpi': dpi,
              'figure.figsize': (args.fig_size[0], args.fig_size[1]),
              'axes.labelsize': font_size,
              'xtick.labelsize': font_size,
              'ytick.labelsize': font_size,
              'figure.titlesize': font_size}
    plt.rcParams.update(params)
    plt.rcParams["font.sans-serif"] = ["Times New Roman"]  # 设置字体
    plt.rcParams["axes.unicode_minus"] = False  # 该语句解决图像中的“-”负号的乱码问题
    plt.style.use('seaborn-whitegrid')
    sns.set_style("white")

    loss = get_loss(args)

    plot_loss_curves(loss,
                     colors=['blue', 'red'],
                     spacing=20,
                     # labels=['train loss', 'test loss'],
                     row_num=len(args.models))
